package com.jsean.demo;

import com.aayushatharva.brotli4j.decoder.BrotliInputStream;
import com.aayushatharva.brotli4j.encoder.BrotliOutputStream;
import com.aayushatharva.brotli4j.encoder.Encoder;
import com.github.luben.zstd.Zstd;
import net.jpountz.lz4.LZ4Compressor;
import net.jpountz.lz4.LZ4Factory;
import net.jpountz.lz4.LZ4FastDecompressor;
import org.xerial.snappy.Snappy;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

public class CompressBenchmark {

    // 测试文件列表
    private static final String[] TEST_FILES = {
            "lorem.txt",
            "json_sample.txt",
            "binary_data.bin"
    };

    // 测试次数
    private static final int WARMUP_ITERATIONS = 3;
    private static final int TEST_ITERATIONS = 5;

    public static void main(String[] args) throws Exception {
        for (String testFile : TEST_FILES) {
            System.out.println("\n===== Testing with file: " + testFile + " =====");

            // 加载测试数据
            byte[] originalData = loadTestData(testFile);
            System.out.printf("Original size: %,d bytes%n", originalData.length);

            // 测试各压缩算法
            testCompressionAlgorithm("LZ4", originalData);
            testCompressionAlgorithm("Zstandard", originalData);
            testCompressionAlgorithm("Snappy", originalData);
            testCompressionAlgorithm("GZIP", originalData);
            testCompressionAlgorithm("Brotli", originalData);
        }
    }

    private static byte[] loadTestData(String filename) throws IOException {
        try (InputStream is = CompressBenchmark.class.getClassLoader()
                .getResourceAsStream(filename);
             ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
            byte[] buffer = new byte[1024];
            int len;
            while ((len = is.read(buffer)) > -1) {
                bos.write(buffer, 0, len);
            }
            return bos.toByteArray();
        }
    }

    private static void testCompressionAlgorithm(String algorithm, byte[] originalData) throws Exception {
        System.out.println("\n--- Testing " + algorithm + " ---");

        // 预热
        for (int i = 0; i < WARMUP_ITERATIONS; i++) {
            runCompressionTest(algorithm, originalData, false);
        }

        // 正式测试
        long totalCompressTime = 0;
        long totalDecompressTime = 0;
        int compressedSize = 0;

        for (int i = 0; i < TEST_ITERATIONS; i++) {
            TestResult result = runCompressionTest(algorithm, originalData, true);
            totalCompressTime += result.compressTime;
            totalDecompressTime += result.decompressTime;
            compressedSize = result.compressedSize;
        }

        // 输出结果
        long avgCompressTime = totalCompressTime / TEST_ITERATIONS;
        long avgDecompressTime = totalDecompressTime / TEST_ITERATIONS;
        double compressionRatio = (double) compressedSize / originalData.length * 100;

        System.out.printf("Avg compress time: %,d ns (%,.2f MB/s)%n",
                avgCompressTime,
                calculateThroughput(originalData.length, avgCompressTime));
        System.out.printf("Avg decompress time: %,d ns (%,.2f MB/s)%n",
                avgDecompressTime,
                calculateThroughput(originalData.length, avgDecompressTime));
        System.out.printf("Compressed size: %,d bytes (%.2f%% of original)%n",
                compressedSize, compressionRatio);
    }

    private static double calculateThroughput(int dataSize, long timeNs) {
        double seconds = timeNs / 1_000_000_000.0;
        double megabytes = dataSize / (1024.0 * 1024.0);
        return megabytes / seconds;
    }

    private static TestResult runCompressionTest(String algorithm, byte[] originalData, boolean collectStats) throws Exception {
        byte[] compressed;
        byte[] decompressed;
        long compressStart, compressEnd, decompressStart, decompressEnd;

        // 压缩
        compressStart = System.nanoTime();
        switch (algorithm) {
            case "LZ4":
                compressed = compressLZ4(originalData);
                break;
            case "Zstandard":
                compressed = compressZstd(originalData);
                break;
            case "Snappy":
                compressed = compressSnappy(originalData);
                break;
            case "GZIP":
                compressed = compressGzip(originalData);
                break;
            case "Brotli":
                compressed = compressBrotli(originalData);
                break;
            default:
                throw new IllegalArgumentException("Unknown algorithm: " + algorithm);
        }
        compressEnd = System.nanoTime();

        // 解压
        decompressStart = System.nanoTime();
        switch (algorithm) {
            case "LZ4":
                decompressed = decompressLZ4(compressed, originalData.length);
                break;
            case "Zstandard":
                decompressed = decompressZstd(compressed, originalData.length);
                break;
            case "Snappy":
                decompressed = decompressSnappy(compressed);
                break;
            case "GZIP":
                decompressed = decompressGzip(compressed);
                break;
            case "Brotli":
                decompressed = decompressBrotli(compressed);
                break;
            default:
                throw new IllegalArgumentException("Unknown algorithm: " + algorithm);
        }
        decompressEnd = System.nanoTime();

        // 验证数据完整性
        if (!Arrays.equals(originalData, decompressed)) {
            throw new RuntimeException("Decompressed data does not match original for " + algorithm);
        }

        return new TestResult(
                compressEnd - compressStart,
                decompressEnd - decompressStart,
                compressed.length
        );
    }

    // LZ4 压缩
    private static byte[] compressLZ4(byte[] data) {
        LZ4Compressor compressor = LZ4Factory.fastestInstance().fastCompressor();
        int maxCompressedLength = compressor.maxCompressedLength(data.length);
        byte[] compressed = new byte[maxCompressedLength];
        int compressedLength = compressor.compress(data, 0, data.length, compressed, 0, maxCompressedLength);
        return Arrays.copyOf(compressed, compressedLength);
    }

    // LZ4 解压
    private static byte[] decompressLZ4(byte[] compressed, int originalLength) {
        LZ4FastDecompressor decompressor = LZ4Factory.fastestInstance().fastDecompressor();
        byte[] restored = new byte[originalLength];
        decompressor.decompress(compressed, 0, restored, 0, originalLength);
        return restored;
    }

    // Zstandard 压缩
    private static byte[] compressZstd(byte[] data) {
        return Zstd.compress(data);
    }

    // Zstandard 解压
    private static byte[] decompressZstd(byte[] compressed, int originalLength) {
        try {
            return Zstd.decompress(compressed, originalLength);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    // Snappy 压缩
    private static byte[] compressSnappy(byte[] data) {
        try {
            return Snappy.compress(data);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    // Snappy 解压
    private static byte[] decompressSnappy(byte[] compressed) {
        try {
            return Snappy.uncompress(compressed);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    // GZIP 压缩
    private static byte[] compressGzip(byte[] data) {
        try (ByteArrayOutputStream bos = new ByteArrayOutputStream();
             GZIPOutputStream gzip = new GZIPOutputStream(bos)) {
            gzip.write(data);
            gzip.finish();
            return bos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    // GZIP 解压
    private static byte[] decompressGzip(byte[] compressed) {
        try (ByteArrayInputStream bis = new ByteArrayInputStream(compressed);
             GZIPInputStream gzip = new GZIPInputStream(bis);
             ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
            byte[] buffer = new byte[1024];
            int len;
            while ((len = gzip.read(buffer)) > 0) {
                bos.write(buffer, 0, len);
            }
            return bos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    // Brotli 压缩
    private static byte[] compressBrotli(byte[] data) {
        try {
            com.aayushatharva.brotli4j.Brotli4jLoader.ensureAvailability();
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            BrotliOutputStream brotli = new BrotliOutputStream(bos, new Encoder.Parameters().setQuality(4));
            brotli.write(data);
            brotli.close(); // 确保调用close()而不是flush()
            return bos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException("Brotli压缩失败", e);
        }
    }

    // Brotli 解压
    private static byte[] decompressBrotli(byte[] compressed) {
        try (ByteArrayInputStream bis = new ByteArrayInputStream(compressed);
             BrotliInputStream brotli = new BrotliInputStream(bis);
             ByteArrayOutputStream bos = new ByteArrayOutputStream()) {
            byte[] buffer = new byte[8192]; // 增大缓冲区
            int len;
            while ((len = brotli.read(buffer)) != -1) { // 明确检查-1
                bos.write(buffer, 0, len);
            }
            return bos.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException("Brotli解压失败: " + e.getMessage(), e);
        }
    }

    private static class TestResult {
        final long compressTime;
        final long decompressTime;
        final int compressedSize;

        TestResult(long compressTime, long decompressTime, int compressedSize) {
            this.compressTime = compressTime;
            this.decompressTime = decompressTime;
            this.compressedSize = compressedSize;
        }
    }
}